#!/usr/bin/python
'''
Created by Colin Gillespie <c.gillespie@ncl.ac.uk>
$LastChangedDate: 2008-04-15 13:39:12 +0100 (Tue, 15 Apr 2008) $

'''

__revision__ = '$Rev: 105 $'

from mrange import mrange
from Factorial import binomialM
import VariableDict as CM
from GlobalVariables import globalVariables
from LogNormalClosure import *

class NormalClosure(object):

    def __init__(self, mom_star = {}, time=False):
        self.mom_star = mom_star
        if time:
            print 'FALSE TIME'
            self.closure_order = globalVariables()['time_closure_order']
            self.variables  = CM.TimeVariables
        else:
            self.closure_order = globalVariables()['closure_order']
            self.variables  = CM.Variables
            

    def getMomStar(self, tup):
        if self.mom_star.has_key(tup):
            return self.mom_star[tup]
        else:
            self.mom_star[tup] = self.__getMomStar(tup)
            return self.mom_star[tup]
    
    def putMomStar(self, tup, value):
        if not self.mom_star.has_key(tup):
            self.mom_star[tup] = value
        return True
    
    
    def closeMoment(self, tup):
        '''Closes a moment by converting to cumulant and back again'''
        if sum(tup) == 0 or sum(tup) == 1:
            raise 'Error: Closure level set too low'
    
        tup_list = list(tup)
        tup_list_rev = list(tup)
        tup_list_rev.reverse()

        for i, value in enumerate(tup_list_rev):
            if value != 0:
                tup_list[-i-1] = tup_list[-i-1] - 1
                break

        value = 0

        for i_tup in mrange(tup_list):
            r_i_diff = tuple([r-i for r, i in zip(tup, i_tup)])
            inner = binomialM(tuple(tup_list), i_tup)
            inner *= self.variables().getMoment(r_i_diff, moment=False)
            inner *= self.variables().getMoment(i_tup)
            value += inner

        return value.expand()


    def closeCumulant(self, value):
        if sum(value) > self.closure_order:
            return 0
        else:
            moment = self.cumulant2Moment(value)
            return moment.expand()

    def cumulant2Moment(self, tup):
        if sum(tup) == 0:
            return 0
        if sum(tup) == 1:
            return self.variables().getMoment(tup)
    
        tup_list = list(tup)
        tup_list_rev = list(tup)
        tup_list_rev.reverse()

        for i, value in enumerate(tup_list_rev):
            if value != 0:
                tup_list[-i-1] = tup_list[-i-1] - 1
                break

        total = 0
        for i_tup in mrange(tup_list):
            r_i_diff = tuple([r-i for r, i in zip(tup, i_tup)])
            inner = binomialM(tuple(tup_list), i_tup)
            inner *= self.variables().getMoment(r_i_diff)
            inner *= self.getMomStar(i_tup)

            total += inner
        return total.expand()     

    def closeMomentAsCumulant(self, tup):
        '''Closes a moment by converting to cumulant and back again'''
    
        if sum(tup) == 0:
            return 1
        if sum(tup) == 1:
            return self.variables().getCumulant(tup, moment=False)
    
        tup_list = list(tup)
        tup_list_rev = list(tup)
        tup_list_rev.reverse()

        for i, value in enumerate(tup_list_rev):
            if value != 0:
                tup_list[-i-1] = tup_list[-i-1] - 1
                break

        value = 0
        for i_tup in mrange(tup_list):
            r_i_diff = tuple([r-i for r, i in zip(tup, i_tup)])
            inner = binomialM(tuple(tup_list), i_tup)
            inner *= self.variables().getCumulant(r_i_diff, moment=False)
            inner *= self.variables().getCumulant(i_tup)
            value += inner

        return value.expand()


    def __getMomStar(self, tup):
        if sum(tup) == 0:
            return 1
        if sum(tup) == 1:
            self.putMomStar(tup, -CM.Variables().getMoment(tup))
            return self.getMomStar(tup)
    
        tup_list = list(tup)
        tup_list_rev = list(tup)
        tup_list_rev.reverse()

        for i, value in enumerate(tup_list_rev):
            if value != 0:
                tup_list[-i-1] = tup_list[-i-1] - 1
                break
    
        total = 0
        for i_tup in mrange(tup_list):
            r_i_diff = tuple([r-i for r, i in zip(tup, i_tup)])
            inner = binomialM(tuple(tup_list), i_tup)
            inner *= -self.variables().getMoment(r_i_diff, moment=False)
            inner *= self.getMomStar(i_tup)
            total += inner
    
        self.putMomStar(tup, total)
        return total.expand()
        
    def initialise(self):
        self.mom_star.clear()
    
class LogNormalClosure(object):

    def __init__(self):
        pass
    
    def closeMoment(self, tup):
    
        if sum(tup) == 0 or sum(tup) == 1:
            raise 'Error: Closure level set too low'

        if tup[0] > 0:
            ir = len(tup)*[0]
            ir[0] = tup[0]
  
            value = CM.Variables().getMoment(tuple(ir))

                
            tup_list = list(tup)
            tup_list[0] = 0
            
            value *= CM.Variables().getMoment(tuple(tup_list))
            print 'Normal:', value, tup
        else:   
            value = convert(tup)
#        print value, tup
        return value



def getMethod(time=False):
    command= globalVariables()['method'] + 'Closure'
    try:
        value = eval(command)
        return value(time = time)
    except NameError:
        raise 'Error: Unknown output method'

if __name__ == '__main__':
    print NormalClosure().closeMoment((4,))
    
